# Copyright Huawei Technologies Co., Ltd. 2025. All rights reserved.
import argparse
from transformers import AutoModel, AutoTokenizer, AutoConfig
from internvl2_utils import get_tokenized_data, get_textvqa_calibration
from msmodelslim.pytorch.llm_ptq.anti_outlier import AntiOutlierConfig, AntiOutlier
from msmodelslim.pytorch.llm_ptq.llm_ptq_tools import Calibrator, QuantConfig


CPU = "cpu"
NPU = "npu"
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, default='')
    parser.add_argument('--calib_images', type=str, default='./textvqa_val')
    parser.add_argument('--calib_num', type=int, default=30, help='random sample calib num')
    parser.add_argument('--save_directory', type=str, default='')
    parser.add_argument('--part_file_size', type=int, default=None)
    parser.add_argument('--w_bit', type=int, default=8)
    parser.add_argument('--a_bit', type=int, default=8)
    parser.add_argument('--device_type', type=str, choices=[CPU, NPU], default=CPU)
    parser.add_argument('--is_8B_model', type=bool, default=True, help='whether to use 8B model')
    args = parser.parse_args()

    # 1.加载模型
    device_map = CPU if args.device_type == CPU else "auto"
    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
    dtype = config.torch_dtype
    model = AutoModel.from_pretrained(
        args.model_path,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
        device_map=device_map,
        use_safetensors=True,
        trust_remote_code=True).eval()
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True, use_fast=False)

    # 2.调用chat接口
    model.forward = model.chat

    # 3.设置回退层
    disable_names = []
    vision_name = []
    if args.is_8B_model:
        llm_name = [
            "language_model.output",
            "mlp1.1",
            "mlp1.3"
        ]
        for i in range(config.vision_config.num_hidden_layers):
            vision_name.extend(
                [
                    f"vision_model.encoder.layers.{i}.mlp.fc2"
                ]
            )
        for i in range(config.llm_config.num_hidden_layers):
            llm_name.extend([
                f"language_model.model.layers.{i}.feed_forward.w2"
            ])
    else:
        llm_name = [
            "language_model.lm_head",
            "mlp1.1",
            "mlp1.3"
        ]
        for i in range(config.vision_config.num_hidden_layers):
            vision_name.extend([
                f"vision_model.encoder.layers.{i}.mlp.fc1",
                f"vision_model.encoder.layers.{i}.mlp.fc2",
                f"vision_model.encoder.layers.{i}.attn.proj",
                f"vision_model.encoder.layers.{i}.attn.qkv",
            ])
        for i in range(config.llm_config.num_hidden_layers):
            llm_name.extend([
                f"language_model.model.layers.{i}.mlp.down_proj"
            ])
    disable_names.extend(vision_name)
    disable_names.extend(llm_name)
    
    # 4.配置校准集
    calibration_dataset = get_textvqa_calibration(args.calib_images, args.calib_num)
    calib_data = get_tokenized_data(tokenizer, calibration_dataset, dtype=dtype)

    # 5.异常值抑制
    anti_config = AntiOutlierConfig(
        a_bit=8,
        w_bit=8,
        anti_method='m2',
        dev_type='npu',
        dev_id=model.device.index
    )
    anti_outlier = AntiOutlier(model, calib_data=calib_data, cfg=anti_config)
    anti_outlier.process()

    # 6.模型量化
    quant_config = QuantConfig(
        w_bit=args.w_bit,
        a_bit=args.a_bit,
        disable_names=disable_names,
        dev_type=args.device_type,
        dev_id=model.device.index,
        act_method=1,
        mm_tensor=False,
    )
    calibrator = Calibrator(model, quant_config, calib_data=calib_data, disable_level='L0')
    calibrator.run()

    # 7.保存权重
    calibrator.save(args.save_directory, save_type=["safe_tensor"], part_file_size=args.part_file_size)
